Source: LucidJS/src/optvis/param/spatial.js

import * as tf from '@tensorflow/tfjs';
import {inverseDecorrelate} from './color.js';

/**
 * Returns randomly initialized input image variables.
 * @param {*} shape 
 * @param {*} sd 
 * @param {*} initVal 
 */
export function pixelImage(shape, sd=undefined, initVal=undefined) {
  if (sd !== undefined && initVal !== undefined) {
    console.log("`pixel_image` received both an initial value and a sd " +
    "argument. Ignoring sd in favor of the supplied initial value.");
  }

  const _sd = sd ? sd : 0.01;
  initVal = initVal ? initVal : tf.randomNormal(
    shape, 0, _sd, 'float32');
  return tf.variable(initVal);
}

/**
 * Returns image-initialized laplacian pyramid as summed image and individual layers.
 * @param {*} shape image shape
 * @param {*} imgData initial image
 * @param {*} nLevels pyramid layers
 * @param {*} invDec inverse decorrelate
 */
export function makeLaplacianPyramidFromImgData(shape, imgData, nLevels=6,
invDec=true){
  const [summedImg, layers] = tf.tidy(() => {
    const imgShape = [1, imgData.height, imgData.width, 4];
    const [b, h, w, ch] = shape;

    let refImg = tf.tensor(Uint8Array.from(imgData.data),
     imgShape, 'float32');
    refImg = tf.image.resizeBilinear(refImg, [w, h]).div(255);
    refImg = refImg.slice([0, 0, 0, 0], [b, h, w, ch]);
    if(invDec){
      refImg = inverseDecorrelate(refImg);
    }

    let currentImg = tf.zeros(shape, 'float32');
    let layers = [];
    for(let i=nLevels-1; i>=0; i--) {
      const k = Math.pow(2, i);
      const wk = Math.floor(w/k);
      const hk = Math.floor(h/k);
      const diff = refImg.sub(currentImg);
      const layer = tf.image.resizeBilinear(diff, [wk, hk]);
      layers.push(layer.dataSync());
      const upsampledLayer = tf.image.resizeBilinear(layer, [w, h]);
      currentImg = currentImg.add(upsampledLayer);
    }
    return [currentImg, layers];
  });
  return [summedImg, layers];
}

/**
 * Returns image-initialized laplacian pyramid as sum function and individual layer variables.
 * @param {*} shape image shape
 * @param {*} imgData initial image
 * @param {*} nLevels pyramid layers
 * @param {*} decorrelate decorrelate colors
 */
export function baseImageLaplacianPyramid(shape, imgData, nLevels=6, decorrelate){
  if(shape.length !== 4 ){
    throw "shape needs batch dimension!";
  }
  const [summedImg, baseImagePyramid] =
  makeLaplacianPyramidFromImgData(shape, imgData, nLevels, decorrelate);
  const [b, w, h, ch]= shape;

  let trainable = [];
  let level;
  for(let i=0; i<nLevels; i++){
    const k = Math.pow(2, nLevels-1-i);
    const wk = Math.floor(w/k);
    const hk = Math.floor(h/k);
    const _shape = [b, wk, hk, ch];
    const downSampledBaseImage =
    tf.tensor(baseImagePyramid[i], _shape);
    level = tf.variable(downSampledBaseImage, true);
    trainable.push(level);
  }
  let f = (xArray) => {
    let pyramid = tf.zeros(shape);
    for(let i=0; i<xArray.length; i++){
      let level = xArray[i];
      level = tf.image.resizeBilinear(level, [w, h]);
      pyramid = pyramid.add(level);
    }
    return pyramid
  }
  return [f, trainable];
}

/**
 * Returns random-initialized laplacian pyramid as sum function and individual layer variables.
 * Contains outcommented code for fading in pyramid layers.
 * @param {*} shape tensor shape
 * @param {*} sd standard deviation
 * @param {*} nLevels pyramid layers
 */
export function laplacianPyramid(shape, sd, nLevels=6){
  if(shape.length !== 4 ){
    throw "shape needs batch dimension!";
  }
  const [b, w, h, ch]= shape;
  const _sd = sd ? sd : 0.01;

  const gaussianBlur = gaussianBlur2D3x3(ch);

  const start = 0;
  let trainable = [];
  let level;
  for(let n=start; n<nLevels+start; n++){
    const k = Math.pow(2, n);
    const wk = Math.floor(w/k);
    const hk = Math.floor(h/k);
    const _shape = [b, wk, hk, ch];
    level = tf.variable(
      tf.randomNormal(_shape, 0, _sd, 'float32'), true);
    trainable.push(level);
  }
  let f = (xArray, weights) => {
    let pyramid = tf.zeros(shape);
    const _weights = weights ? weights : ones(nLevels);

    for(let i=0; i<xArray.length; i++){
      let level = xArray[i];
      const wght = Math.max(_weights[i], 0.000001);
      if(i > 0 && false) {
        level = gaussianBlur.apply(level);
        level = gaussianBlur.apply(level);
      }
      level = tf.image.resizeBilinear(level, [w, h], true);
      //pyramid = pyramid.add(level.mul(wght));
      pyramid = pyramid.add(level);
    }
    return pyramid
  }
  return [f, trainable];
}

function gaussianBlur2D3x3(ch) {
  const conv2d = tf.layers.depthwiseConv2d(
    {kernelSize: 3, padding: 'same', trainable: false, updatable: false,
      weights: gaussianKernel3x3(ch)});
  return conv2d;
}

/**
 * Computes gaussian 3x3 kernel
 * @param {*} ch channels
 */
function gaussianKernel3x3(ch) {
  const c = 1/16;
  const s = 1/8;
  const m = 1/4;

  let weights = tf.tensor2d([[c, s, c], [s, m, s], [c, s, c]]);
  weights = tf.reshape(weights, [1, 1, 3, 3]);
  weights = tf.transpose(weights, [2, 3, 0, 1]);
  weights = tf.tile(weights, [1, 1, ch, 1]);

  const bias = tf.zeros([ch]);

  return [weights, bias];
}

function ones(n){
  const ret = [];
  for(let i=0; i<n; i++){
    ret.push(1);
  }
  return ret;
}